/*
 *  Copyright 2008-2009 NVIDIA Corporation
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

/*! \file coo_matrix.h
 *  \brief Coordinate matrix format
 */

#pragma once

#include <thrust/set_operations.h>
#include <thrust/unique.h>

#include <cusp/detail/config.h>

#include <cusp/array1d.h>
#include <cusp/format.h>
#include <cusp/coo_matrix.h>
#include <cusp/csr_matrix.h>
#include <cusp/dia_matrix.h>
#include <cusp/ell_matrix.h>
#include <cusp/hyb_matrix.h>

#include <cusp/verify.h>
#include <cusp/print.h>
#include <cusp/detail/matrix_base.h>

#include <cusp/graph/partition.h>
#include <cusp/io/matrix_market.h>

namespace cusp
{

struct mult_method_type {};

struct multi_sync_mult     : mult_method_type {};
struct multi_host_mult     : mult_method_type {};
struct multi_kernel_mult   : mult_method_type {};
struct multi_p2p_copy_mult : mult_method_type {};

template <typename T>
struct equal_to
{
    const T num;

    equal_to(const T num) : num(num) {}

    __host__ __device__ bool operator()(const T &x) const {
        return x == num;
    }
};

template <typename IndexType>
struct is_local_nnz
{
    typedef IndexType argument_type;
    typedef bool result_type;
    typedef bool type;

    const   IndexType lowest_row;
    const   IndexType highest_row;

    is_local_nnz(const IndexType lowest_row, const IndexType highest_row)
        : lowest_row(lowest_row), highest_row(highest_row) {}

    __host__ __device__
    bool operator()(const IndexType column_index) const
    {
        return (column_index >= lowest_row) && (column_index <= highest_row);
    }
};

template< typename IndexType, typename ValueType, typename MultMethodType >
struct RemoteContainerType {};

template< typename IndexType, typename ValueType >
struct RemoteContainerType<IndexType,ValueType,multi_sync_mult>
{
    typedef typename cusp::coo_matrix<IndexType,ValueType,cusp::device_memory> type;
};
template< typename IndexType, typename ValueType >
struct RemoteContainerType<IndexType,ValueType,multi_host_mult>
{
    typedef typename cusp::csr_matrix<IndexType,ValueType,cusp::host_memory> type;
};
template< typename IndexType, typename ValueType >
struct RemoteContainerType<IndexType,ValueType,multi_kernel_mult>
{
    typedef typename cusp::coo_matrix<IndexType,ValueType,cusp::device_memory> type;
};
template< typename IndexType, typename ValueType >
struct RemoteContainerType<IndexType,ValueType,multi_p2p_copy_mult>
{
    typedef typename cusp::coo_matrix<IndexType,ValueType,cusp::device_memory> type;
};


template< typename IndexType, typename ValueType, typename MemorySpace, typename Format >
struct SingleContainerType {};

template< typename IndexType, typename ValueType, typename MemorySpace >
struct SingleContainerType<IndexType,ValueType,MemorySpace,cusp::multi_coo_format>
{
    typedef typename cusp::coo_matrix<IndexType,ValueType,MemorySpace> type;
};

template< typename IndexType, typename ValueType, typename MemorySpace >
struct SingleContainerType<IndexType,ValueType,MemorySpace,cusp::multi_csr_format>
{
    typedef typename cusp::csr_matrix<IndexType,ValueType,MemorySpace> type;
};

template< typename IndexType, typename ValueType, typename MemorySpace >
struct SingleContainerType<IndexType,ValueType,MemorySpace,cusp::multi_dia_format>
{
    typedef typename cusp::dia_matrix<IndexType,ValueType,MemorySpace> type;
};

template< typename IndexType, typename ValueType, typename MemorySpace >
struct SingleContainerType<IndexType,ValueType,MemorySpace,cusp::multi_ell_format>
{
    typedef typename cusp::ell_matrix<IndexType,ValueType,MemorySpace> type;
};

template< typename IndexType, typename ValueType, typename MemorySpace >
struct SingleContainerType<IndexType,ValueType,MemorySpace,cusp::multi_hyb_format>
{
    typedef typename cusp::hyb_matrix<IndexType,ValueType,MemorySpace> type;
};


/*! \addtogroup sparse_matrices Sparse Matrices
 */

/*! \addtogroup sparse_matrix_containers Sparse Matrix Containers
 *  \ingroup sparse_matrices
 *  \{
 */

/*! \p coo_matrix : Coordinate matrix container
 *
 * \tparam IndexType Type used for matrix indices (e.g. \c int).
 * \tparam ValueType Type used for matrix values (e.g. \c float).
 * \tparam MemorySpace A memory space (e.g. \c cusp::host_memory or cusp::device_memory)
 *
 * \note The matrix entries must be sorted by row index.
 * \note The matrix should not contain duplicate entries.
 *
 *  The following code snippet demonstrates how to create a 4-by-3
 *  \p coo_matrix on the host with 6 nonzeros and then copies the
 *  matrix to the device.
 *
 *  \code
 *  #include <cusp/coo_matrix.h>
 *  ...
 *
 *  // allocate storage for (4,3) matrix with 6 nonzeros
 *  cusp::coo_matrix<int,float,cusp::host_memory> A(4,3,6);
 *
 *  // initialize matrix entries on host
 *  A.row_indices[0] = 0; A.column_indices[0] = 0; A.values[0] = 10;
 *  A.row_indices[1] = 0; A.column_indices[1] = 2; A.values[1] = 20;
 *  A.row_indices[2] = 2; A.column_indices[2] = 2; A.values[2] = 30;
 *  A.row_indices[3] = 3; A.column_indices[3] = 0; A.values[3] = 40;
 *  A.row_indices[4] = 3; A.column_indices[4] = 1; A.values[4] = 50;
 *  A.row_indices[5] = 3; A.column_indices[5] = 2; A.values[5] = 60;
 *
 *  // A now represents the following matrix
 *  //    [10  0 20]
 *  //    [ 0  0  0]
 *  //    [ 0  0 30]
 *  //    [40 50 60]
 *
 *  // copy to the device
 *  cusp::coo_matrix<int,float,cusp::device_memory> B = A;
 *  \endcode
 *
 */

namespace detail
{
template< typename PermutationMatrix, typename Matrix >
void permuteMatrix( const PermutationMatrix& P, Matrix& coo, cusp::graph::rowwise )
{
    typedef typename Matrix::index_type IndexType;
    typedef typename Matrix::value_type ValueType;
    typedef typename Matrix::memory_space MemorySpace;

    cusp::csr_matrix<IndexType,ValueType,cusp::device_memory> csr(coo);
    bool is_symmetric = true; //cusp::graph::is_structurally_symmetric(csr);

    cusp::coo_matrix<IndexType, ValueType, cusp::device_memory> coo_d(coo);
    cusp::coo_matrix<IndexType, ValueType, cusp::device_memory> cooP(coo_d);
    cusp::permutation_matrix<int,cusp::device_memory> P_d(P);
    if( is_symmetric ) {
        cusp::multiply(cooP,P_d,coo_d);
        cooP = coo_d;
    }
    cusp::multiply(P_d,cooP,coo_d);
    coo = coo_d;
}

template< typename PermutationMatrix, typename Matrix >
void permuteMatrix( const PermutationMatrix& P, Matrix& coo, cusp::graph::colwise )
{
    typedef typename Matrix::index_type IndexType;
    typedef typename Matrix::value_type ValueType;
    typedef typename Matrix::memory_space MemorySpace;

    cusp::csr_matrix<IndexType,ValueType,cusp::device_memory> csr(coo);
    bool is_symmetric = cusp::graph::is_structurally_symmetric(csr);

    cusp::coo_matrix<IndexType, ValueType, cusp::device_memory> coo_d(coo);
    cusp::coo_matrix<IndexType, ValueType, cusp::device_memory> cooP(coo_d);
    cusp::permutation_matrix<int,cusp::device_memory> P_d(P);
    cusp::multiply(cooP,P_d,coo_d);
    if( is_symmetric ) {
        cooP = coo_d;
        cusp::multiply(P_d,cooP,coo_d);
    }
    coo = coo_d;
}

template< typename MatrixType >
size_t active_warps( const MatrixType& A )
{
    typedef typename MatrixType::index_type IndexType;
    typedef typename MatrixType::value_type ValueType;

    const unsigned int BLOCK_SIZE = 256;
    const unsigned int MAX_BLOCKS = cusp::detail::device::arch::max_active_blocks(
                                        cusp::detail::device::spmv_coo_flat_kernel<IndexType, ValueType, BLOCK_SIZE, false>, BLOCK_SIZE, (size_t) 0);
    const unsigned int WARPS_PER_BLOCK = BLOCK_SIZE / WARP_SIZE;

    const unsigned int num_units  = A.num_entries / WARP_SIZE;
    const unsigned int num_warps  = std::min(num_units, WARPS_PER_BLOCK * MAX_BLOCKS);
    const unsigned int num_iters  = DIVIDE_INTO(num_units, num_warps);

    const unsigned int interval_size = WARP_SIZE * num_iters;

    const IndexType tail = num_units * WARP_SIZE; // do the last few nonzeros separately (fewer than WARP_SIZE elements)

    return (interval_size == 0) ? 0 : DIVIDE_INTO(tail, interval_size);
}

} // end detail namespace

template <typename IndexType, typename ValueType, class MemorySpace, typename Format, typename MultMethodType, typename PartitionerType, typename MethodType>
class multi_matrix : public detail::matrix_base<IndexType,ValueType,MemorySpace,Format>
{
    typedef cusp::detail::matrix_base<IndexType,ValueType,MemorySpace,Format> Parent;

    template <typename MatrixType>
    void CreateMultiMatrix(const MatrixType& matrix)
    {
        if( partition == NULL )
            throw cusp::invalid_input_exception("partition is set to NULL");

        if ( cusp::is_valid_partition(*partition) )
        {
            // create permutation matrix from partition
            cusp::permutation_matrix<int,cusp::host_memory> P;
            cusp::graph::partitionToPermMatrix(*partition, P);

            // permute the matrix
            cusp::coo_matrix<IndexType, ValueType, cusp::host_memory> coo(matrix);
            coo.sort_by_row_and_column();
            detail::permuteMatrix( P, coo, MethodType() );

            // identify unique devices in partition
            int num_slices;
            {
                cusp::array1d<int,cusp::host_memory> unique( *partition );
                thrust::sort(unique.begin(), unique.end());
                num_slices = thrust::unique( unique.begin(), unique.end() ) - unique.begin();
                unique.resize(num_slices);
                deviceNum = std::vector<size_t>(unique.begin(), unique.end());
            }

            // initialize all arrays to correct length
            matrix_array.resize(num_slices);
            local_rows.resize(num_slices);
            real_local_rows.resize(num_slices);
            temp_rows1.resize(num_slices);
            temp_rows2.resize(num_slices);
            temp_vals1.resize(num_slices);
            temp_vals2.resize(num_slices);
            temp_y.resize(num_slices);

            cusp::csr_matrix<IndexType,ValueType,cusp::device_memory> csr(coo);
            bool is_symmetric = true; //cusp::graph::is_structurally_symmetric(csr);
    	    //printf("Is symmetric : %s\n", is_symmetric ? "Yes" : "No");

            if( is_symmetric )
            {
                if(  thrust::detail::is_same<MethodType, cusp::graph::rowwise>::value )
                    if( thrust::detail::is_same<MultMethodType, cusp::multi_p2p_copy_mult>::value ||
                            thrust::detail::is_same<MultMethodType, cusp::multi_host_mult>::value ) {
                        unique_remote_rows.resize(num_slices);
                        remote_rows.resize(num_slices);
                        remote_array.resize(num_slices);
                        copy_remote_rows.resize(num_slices);
                        gather_remote_rows.resize(num_slices);
                    }
            }

            // count the number of rows in each partition
            std::vector<size_t> partition_row_sizes(num_slices+1, 0);
            for ( IndexType index = 1; index < num_slices+1; index++ )
                partition_row_sizes[index] = thrust::count( partition->begin(), partition->end(), int(deviceNum[index-1]) ) + partition_row_sizes[index-1];

            if( thrust::detail::is_same<MethodType, cusp::graph::colwise>::value ) {
                cusp::coo_matrix<IndexType,ValueType,cusp::host_memory> coo_(coo);
                cusp::transpose(coo_, coo);
                coo.sort_by_row_and_column();
            }

            cusp::array1d<IndexType,cusp::host_memory> row_offsets(coo.num_rows+1);
            cusp::detail::indices_to_offsets(coo.row_indices, row_offsets);

            // count the number of nonzero entries in each partition
            std::vector<size_t> partition_entry_sizes(num_slices+1, 0);
            for ( IndexType index = 1; index < num_slices+1; index++ )
                partition_entry_sizes[index] = row_offsets[partition_row_sizes[index]];

            for ( IndexType index = 0; index < num_slices; index++ )
            {
                cudaSetDevice( deviceNum[index] );

                int lowest_row = partition_row_sizes[index];
                int highest_row = partition_row_sizes[index+1]-1;

                size_t num_coo_rows = highest_row - lowest_row + 1;
                size_t num_coo_cols = coo.num_cols;
                size_t num_coo_entries = partition_entry_sizes[index+1] - partition_entry_sizes[index];

                int start = partition_entry_sizes[index];
                int end   = partition_entry_sizes[index+1];

                cusp::coo_matrix<IndexType, ValueType, cusp::host_memory> local_coo =
                    cusp::make_coo_matrix_view( num_coo_rows, num_coo_cols, num_coo_entries,
                                                cusp::make_array1d_view( coo.row_indices.begin()+start   , coo.row_indices.begin()+end ),
                                                cusp::make_array1d_view( coo.column_indices.begin()+start, coo.column_indices.begin()+end ),
                                                cusp::make_array1d_view( coo.values.begin()+start        , coo.values.begin()+end ) );
                coo.sort_by_row_and_column();

                typedef typename cusp::array1d_view< thrust::constant_iterator<IndexType> > ConstantView;
                ConstantView constant_view( thrust::constant_iterator<IndexType>(lowest_row),
                                            thrust::constant_iterator<IndexType>(lowest_row)+num_coo_entries );
                cusp::blas::axpy( constant_view, local_coo.row_indices, IndexType(-1) );

                if( thrust::detail::is_same<MethodType, cusp::graph::rowwise>::value )
                {
                    if( is_symmetric )
                    {
                        if( thrust::detail::is_same<MultMethodType, cusp::multi_p2p_copy_mult>::value ||
                                thrust::detail::is_same<MultMethodType, cusp::multi_host_mult>::value )
                        {
                            int their_index = (index+1)%num_slices;

                            size_t num_local_nnzs = thrust::count_if( local_coo.column_indices.begin(), local_coo.column_indices.end(),
                                                    is_local_nnz<IndexType>(lowest_row, highest_row));
                            cusp::coo_matrix<IndexType, ValueType, cusp::host_memory> local_clean_coo( num_coo_rows, num_coo_cols, num_local_nnzs );

                            thrust::copy_if
                            (thrust::make_zip_iterator(thrust::make_tuple(local_coo.row_indices.begin(), local_coo.column_indices.begin(), local_coo.values.begin())),
                             thrust::make_zip_iterator(thrust::make_tuple(local_coo.row_indices.end()  , local_coo.column_indices.end()  , local_coo.values.end()  )),
                             local_coo.column_indices.begin(),
                             thrust::make_zip_iterator(thrust::make_tuple(local_clean_coo.row_indices.begin(), local_clean_coo.column_indices.begin(), local_clean_coo.values.begin())),
                             is_local_nnz<IndexType>(lowest_row, highest_row));

                            ConstantView constant_view_local( thrust::constant_iterator<IndexType>(lowest_row),
                                                              thrust::constant_iterator<IndexType>(lowest_row)+num_local_nnzs );
                            cusp::blas::axpy( constant_view_local, local_clean_coo.column_indices, IndexType(-1) );

                            size_t num_remote_nnzs = local_coo.num_entries - num_local_nnzs;
                            cusp::coo_matrix<IndexType, ValueType, cusp::host_memory> remote_coo( num_coo_rows, num_coo_cols, num_remote_nnzs );

                            thrust::copy_if
                            (thrust::make_zip_iterator(thrust::make_tuple(local_coo.row_indices.begin(), local_coo.column_indices.begin(), local_coo.values.begin())),
                             thrust::make_zip_iterator(thrust::make_tuple(local_coo.row_indices.end()  , local_coo.column_indices.end()  , local_coo.values.end()  )),
                             local_coo.column_indices.begin(),
                             thrust::make_zip_iterator(thrust::make_tuple(remote_coo.row_indices.begin(), remote_coo.column_indices.begin(), remote_coo.values.begin())),
                             thrust::not1(is_local_nnz<IndexType>(lowest_row, highest_row)));

                            cusp::array1d<IndexType,cusp::host_memory> remote_column_indices(remote_coo.column_indices);
                            cusp::array1d<IndexType,cusp::host_memory> new_remote_indices(coo.num_cols,-1);
                            thrust::sort(remote_column_indices.begin(), remote_column_indices.end());
                            int num_remote_column_indices = thrust::unique(remote_column_indices.begin(), remote_column_indices.end()) - remote_column_indices.begin();
                            remote_column_indices.resize(num_remote_column_indices);
                            thrust::scatter(thrust::counting_iterator<IndexType>(0),
                                            thrust::counting_iterator<IndexType>(num_remote_column_indices),
                                            remote_column_indices.begin(), new_remote_indices.begin());
                            thrust::gather( remote_coo.column_indices.begin(), remote_coo.column_indices.end(),
                                            new_remote_indices.begin(),
                                            remote_coo.column_indices.begin() );

                            int shift_value = (highest_row+1)%coo.num_rows;
                            ConstantView constant_remote_shift( thrust::constant_iterator<IndexType>(shift_value),
                                                                thrust::constant_iterator<IndexType>(shift_value)+num_remote_column_indices );
                            cusp::blas::axpy( constant_remote_shift, remote_column_indices, IndexType(-1) );

                            cusp::blas::fill( new_remote_indices, -1 );
                            cusp::array1d<IndexType,cusp::host_memory> remote_row_indices(remote_coo.row_indices);
                            thrust::sort(remote_row_indices.begin(), remote_row_indices.end());
                            int num_remote_row_indices = thrust::unique(remote_row_indices.begin(), remote_row_indices.end()) - remote_row_indices.begin();
                            remote_row_indices.resize(num_remote_row_indices);
                            thrust::scatter(thrust::counting_iterator<IndexType>(0),
                                            thrust::counting_iterator<IndexType>(num_remote_row_indices),
                                            remote_row_indices.begin(), new_remote_indices.begin());
                            thrust::gather( remote_coo.row_indices.begin(), remote_coo.row_indices.end(),
                                            new_remote_indices.begin(),
                                            remote_coo.row_indices.begin() );

                            remote_coo.resize( num_remote_row_indices, num_remote_column_indices, num_remote_nnzs );

                            local_rows[index].resize(local_clean_coo.num_rows);
                            real_local_rows[index].resize(local_clean_coo.num_rows);
                            remote_rows[index].resize(num_slices);
                            copy_remote_rows[index].resize(num_slices);
                            gather_remote_rows[index].resize(num_slices);

                            thrust::sequence(local_rows[index].begin(), local_rows[index].end(), 0);
                            thrust::sequence(real_local_rows[index].begin(), real_local_rows[index].end(), lowest_row);

                            remote_array[index] = remote_coo;
                            unique_remote_rows[index] = remote_row_indices;
                            remote_rows[index][their_index] = remote_column_indices;
                            copy_remote_rows[index][their_index].resize(remote_row_indices.size());
                            gather_remote_rows[index][their_index].resize(remote_row_indices.size());

                            local_rows[index].setDeviceNum( deviceNum[index] );
                            real_local_rows[index].setDeviceNum( deviceNum[index] );
                            remote_array[index].setDeviceNum( deviceNum[index] );
                            unique_remote_rows[index].setDeviceNum( deviceNum[index] );
                            remote_rows[index][their_index].setDeviceNum( deviceNum[index] );
                            copy_remote_rows[index][their_index].setDeviceNum( deviceNum[index] );
                            gather_remote_rows[index][their_index].setDeviceNum( deviceNum[index] );

                            try
                            {
                                matrix_array[index] = local_clean_coo;
                                matrix_array[index].setDeviceNum( deviceNum[index] );
                            }
                            catch (cusp::format_conversion_exception)
                            {
                                matrix_array[index].setDeviceNum( deviceNum[index] );
				throw;
                            }
                        }
                        else
                        {
                            local_rows[index].resize(local_coo.num_rows);
                            real_local_rows[index].resize(local_coo.num_rows);

                            thrust::sequence(local_rows[index].begin(), local_rows[index].end(), 0);
                            thrust::sequence(real_local_rows[index].begin(), real_local_rows[index].end(), lowest_row);

                            local_rows[index].setDeviceNum( deviceNum[index] );
                            real_local_rows[index].setDeviceNum( deviceNum[index] );

                            try
                            {
                                matrix_array[index] = local_coo;
                                matrix_array[index].setDeviceNum( deviceNum[index] );
                            }
                            catch (cusp::format_conversion_exception)
                            {
                                matrix_array[index].setDeviceNum( deviceNum[index] );
				throw;
                            }
                        }
                    }
                    else
                    {
                        local_rows[index].resize(local_coo.num_rows);
                        real_local_rows[index].resize(local_coo.num_rows);

                        thrust::sequence(local_rows[index].begin(), local_rows[index].end(), 0);
                        thrust::sequence(real_local_rows[index].begin(), real_local_rows[index].end(), lowest_row);

                        local_rows[index].setDeviceNum( deviceNum[index] );
                        real_local_rows[index].setDeviceNum( deviceNum[index] );

                        try
                        {
                            matrix_array[index] = local_coo;
                            matrix_array[index].setDeviceNum( deviceNum[index] );
                        }
                        catch (cusp::format_conversion_exception)
                        {
                            matrix_array[index].setDeviceNum( deviceNum[index] );
			    throw;
                        }
                    }
                }
                else if(thrust::detail::is_same<MethodType, cusp::graph::colwise>::value )
                {
                    local_rows[index].resize(local_coo.num_rows);
                    real_local_rows[index].resize(local_coo.num_rows);

                    cusp::coo_matrix<IndexType,ValueType,cusp::host_memory> local_coo_(local_coo);
                    cusp::transpose(local_coo_, local_coo);

                    thrust::sequence(local_rows[index].begin(), local_rows[index].end(), 0);
                    thrust::sequence(real_local_rows[index].begin(), real_local_rows[index].end(), lowest_row);
                    local_rows[index].setDeviceNum( deviceNum[index] );
                    real_local_rows[index].setDeviceNum( deviceNum[index] );

                    matrix_array[index] = local_coo;
                    matrix_array[index].setDeviceNum( deviceNum[index] );
                }
                else
                {
                    throw cusp::invalid_input_exception("Unknown or currently unsupported partition method specified");
                }

                size_t temp_size = 4600; //detail::active_warps( local_coo );
                temp_rows1[index] = cusp::array1d<IndexType,cusp::host_memory>(temp_size, 0);
                temp_rows2[index] = cusp::array1d<IndexType,cusp::host_memory>(temp_size, 0);
                temp_vals1[index] = cusp::array1d<ValueType,cusp::host_memory>(temp_size, 0);
                temp_vals2[index] = cusp::array1d<ValueType,cusp::host_memory>(temp_size, 0);
                temp_y[index]     = cusp::array1d<ValueType,cusp::host_memory>( remote_array.size() > 0 ? remote_array[index].num_rows : local_coo.num_rows, 0);

                temp_rows1[index].setDeviceNum( deviceNum[index] );
                temp_rows2[index].setDeviceNum( deviceNum[index] );
                temp_vals1[index].setDeviceNum( deviceNum[index] );
                temp_vals2[index].setDeviceNum( deviceNum[index] );
                temp_y[index].setDeviceNum( deviceNum[index] );
            }

            if( thrust::detail::is_same<MethodType, cusp::graph::rowwise>::value )
            {
                if( is_symmetric )
                {
                    if( thrust::detail::is_same<MultMethodType, cusp::multi_host_mult>::value )
                    {
                        cudaSetDevice(deviceNum[0]);
                        cusp::array1d<IndexType,cusp::host_memory> temp1 = remote_rows[0][1];
                        cudaSetDevice(deviceNum[1]);
                        cusp::array1d<IndexType,cusp::host_memory> temp2 = remote_rows[1][0];
                        remote_rows[1][0].resize(temp1.size());
                        remote_rows[1][0] = temp1;
                        cudaSetDevice(deviceNum[0]);
                        remote_rows[0][1].resize(temp2.size());
                        remote_rows[0][1] = temp2;
                    }
                }
            }
        }
    }
public:

    cusp::array1d<int,cusp::host_memory> private_partition;

    /*! index array container type
     */
    typedef typename cusp::array1d<IndexType,MemorySpace> IndexArray;
    /*! value array container type
     */
    typedef typename cusp::array1d<ValueType,MemorySpace> ValueArray;

    /*! equivalent container type
     */
    typedef typename cusp::multi_matrix<IndexType, ValueType, MemorySpace, Format, MultMethodType, PartitionerType, MethodType> container;

    /*! equivalent single GPU container type
     */
    typedef typename SingleContainerType<IndexType,ValueType,MemorySpace,Format>::type SingleContainer;
    /*! equivalent remote nonzeros container type
     */
    typedef typename RemoteContainerType<IndexType,ValueType,MultMethodType>::type RemoteContainer;
    /*! gather value array container type
     */
    typedef typename RemoteContainer::memory_space CopyMemorySpace;
    typedef typename cusp::array1d<ValueType,CopyMemorySpace> CopyValueArray;

    std::vector< size_t > 	   	   deviceNum;          	// list of device numbers for each item
    std::vector< IndexArray > 		   local_rows;		// local rows(columns) of the matrix owned by each entry
    std::vector< IndexArray > 		   real_local_rows;	// local rows(columns) of the matrix owned by each entry
    std::vector< IndexArray > 		   unique_remote_rows;	// unique set of rows referenced in remote matrix from say GPU0 to GPU1
    std::vector< IndexArray > 		   temp_rows1;		// temporary row memory for COO multiplication of remote array
    std::vector< IndexArray > 		   temp_rows2;		// temporary row memory for COO multiplication of COO or HYB matrix
    std::vector< ValueArray > 		   temp_vals1;		// temporary value memory for COO multiplication
    std::vector< ValueArray > 		   temp_vals2;		// temporary value memory for COO multiplication
    std::vector< ValueArray > 		   temp_y;		// temporary storage for reduction used in P2P and Host kernels
    std::vector< SingleContainer > 	   matrix_array;	// array of (local) matrices of SingleContainer type
    std::vector< RemoteContainer > 	   remote_array;	// array of (remote)matrices of RemoteContainer type
    std::vector< std::vector<IndexArray> > remote_rows;		// set of rows referenced by remote matrix between each pair of devices
    //std::vector< std::vector<ValueArray> > copy_remote_rows;	// storage space for copying remote entries
    std::vector< std::vector<CopyValueArray> > copy_remote_rows;// storage space for copying remote entries
    std::vector< std::vector<ValueArray> > gather_remote_rows;  // storage space for gathering remote entries before copying

    const cusp::array1d<int,cusp::host_memory> * partition;

    multi_matrix(void) : Parent(0, 0, 0), partition(NULL) {}

    multi_matrix(const multi_matrix& other_matrix) : Parent(other_matrix.num_rows, other_matrix.num_cols, other_matrix.num_entries)
    {
        deviceNum = other_matrix.deviceNum;
        private_partition = other_matrix.private_partition;
        partition = &private_partition;

        int num_slices = other_matrix.getNumSlices();

        matrix_array.resize( num_slices );
        local_rows.resize( num_slices );
        real_local_rows.resize( num_slices );
        temp_rows1.resize(num_slices);
        temp_rows2.resize(num_slices);
        temp_vals1.resize(num_slices);
        temp_vals2.resize(num_slices);
        temp_y.resize(num_slices);

        bool is_symmetric = (other_matrix.remote_array.size() > 0);

        if( is_symmetric )
        {
            if(  thrust::detail::is_same<MethodType, cusp::graph::rowwise>::value )
                if( thrust::detail::is_same<MultMethodType, cusp::multi_p2p_copy_mult>::value ||
                        thrust::detail::is_same<MultMethodType, cusp::multi_host_mult>::value ) {
                    remote_rows.resize(num_slices);
                    remote_array.resize(num_slices);
                    unique_remote_rows.resize(num_slices);
                    copy_remote_rows.resize(num_slices);
                    gather_remote_rows.resize(num_slices);
                }
        }

        for( int index = 0; index < num_slices; index++ )
        {
            cudaSetDevice(deviceNum[index]);

            matrix_array[index] = other_matrix(index);
            local_rows[index]   = other_matrix.local_rows[index];
            real_local_rows[index]= other_matrix.local_rows[index];
            temp_rows1[index]	= other_matrix.temp_rows1[index];
            temp_rows2[index]	= other_matrix.temp_rows2[index];
            temp_vals1[index]	= other_matrix.temp_vals1[index];
            temp_vals2[index]	= other_matrix.temp_vals2[index];
            temp_y[index]	= other_matrix.temp_y[index];

            matrix_array[index].setDeviceNum(deviceNum[index]);
            local_rows[index].setDeviceNum(deviceNum[index]);
            real_local_rows[index].setDeviceNum(deviceNum[index]);
            temp_rows1[index].setDeviceNum(deviceNum[index]);
            temp_rows2[index].setDeviceNum(deviceNum[index]);
            temp_vals1[index].setDeviceNum(deviceNum[index]);
            temp_vals2[index].setDeviceNum(deviceNum[index]);
            temp_y[index].setDeviceNum(deviceNum[index]);

            if( remote_array.size() > 0 )
            {
                remote_array[index]       = other_matrix.remote_array[index];
                unique_remote_rows[index] = other_matrix.unique_remote_rows[index];

                remote_array[index].setDeviceNum(deviceNum[index]);
                unique_remote_rows[index].setDeviceNum(deviceNum[index]);

                remote_rows[index].resize( other_matrix.remote_rows[index].size() );
                copy_remote_rows[index].resize( other_matrix.copy_remote_rows[index].size() );
                gather_remote_rows[index].resize( other_matrix.gather_remote_rows[index].size() );

                for( int index1 = 0; index1 < int(other_matrix.remote_rows[index].size()); index1++ ) {
                    remote_rows[index][index1] = other_matrix.remote_rows[index][index1];
                    copy_remote_rows[index][index1] = other_matrix.copy_remote_rows[index][index1];
                    gather_remote_rows[index][index1] = other_matrix.gather_remote_rows[index][index1];

                    remote_rows[index][index1].setDeviceNum(deviceNum[index]);
                    copy_remote_rows[index][index1].setDeviceNum(deviceNum[index]);
                    gather_remote_rows[index][index1].setDeviceNum(deviceNum[index]);
                }
            }
        }
    }

    /*! Construct a \p coo_matrix with a specific shape and number of nonzero entries.
     *
     *  \param num_rows Number of rows.
     *  \param num_cols Number of columns.
     *  \param num_entries Number of nonzero matrix entries.
     */
    template <typename ArrayType>
    multi_matrix(size_t num_rows, size_t num_cols, size_t num_entries, const ArrayType& part)
        : Parent(num_rows, num_cols, num_entries),
          partition(&part) {}

    /*! Construct a \p coo_matrix from another matrix.
     *
     *  \param matrix Another sparse or dense matrix.
     */
    template <typename MatrixType>
    multi_matrix(const MatrixType& matrix)
        : Parent(matrix.num_rows, matrix.num_cols, matrix.num_entries)
    {
        int deviceCount;
        cudaGetDeviceCount(&deviceCount);

        cusp::graph::partition<PartitionerType,MethodType>(matrix, deviceCount, private_partition);
        partition = &private_partition;

        CreateMultiMatrix(matrix);
    }

    template <typename MatrixType, typename ArrayType>
    multi_matrix(const MatrixType& matrix, const ArrayType& part)
        : Parent(matrix.num_rows, matrix.num_cols, matrix.num_entries),
          partition(&part) {
        CreateMultiMatrix(matrix);
    }

    template <typename MatrixType, typename ArrayType>
    multi_matrix(const MatrixType& matrix, const ArrayType* part)
        : Parent(matrix.num_rows, matrix.num_cols, matrix.num_entries),
          partition(part) {
        CreateMultiMatrix(matrix);
    }

    SingleContainer &operator()(const size_t index) {
        return matrix_array[index];
    }

    const SingleContainer &operator()(const size_t index) const {
        return matrix_array[index];
    }

    size_t getNumSlices(void) const {
        return matrix_array.size();
    }

    int getDeviceNum( int index ) const {
        return deviceNum[index];
    }

    size_t numRemoteNonzeros(void)
    {
        size_t total_nnzs = 0;
        for( int index = 0; index < remote_array.size(); index++ )
            total_nnzs += remote_array[index].num_entries;

        return total_nnzs;
    }

    size_t numCuts(void)
    {
        size_t total_cuts = 0;
        for( int index1 = 0; index1 < remote_rows.size(); index1++ )
            for( int index2 = 0; index2 < remote_rows[index1].size(); index2++ )
                total_cuts += remote_rows[index1][index2].size();

        return total_cuts;
    }

}; // class coo_matrix
/*! \}
 */

template <typename IndexType, typename ValueType, class MemorySpace, typename MultMethodType,
         typename PartitionerType = cusp::graph::patoh, typename MethodType = cusp::graph::rowwise>
class multi_coo_matrix : public multi_matrix<IndexType,ValueType,MemorySpace,cusp::multi_coo_format,MultMethodType,PartitionerType,MethodType>
{
    typedef multi_matrix<IndexType,ValueType,MemorySpace,cusp::multi_coo_format,MultMethodType,PartitionerType,MethodType> Parent;

public:
    typedef MultMethodType multi_mult_method;
    typedef MethodType part_method;

    multi_coo_matrix(void) : Parent() {}

    template <typename ArrayType>
    multi_coo_matrix(size_t num_rows, size_t num_cols, size_t num_entries, const ArrayType& part)
        : Parent(num_rows, num_cols, num_entries, num_entries, part) {}

    template <typename MatrixType>
    multi_coo_matrix(const MatrixType& matrix)
        : Parent(matrix) {}

    template <typename MatrixType, typename ArrayType>
    multi_coo_matrix(const MatrixType& matrix, const ArrayType& part)
        : Parent(matrix, part) {}

    multi_coo_matrix(const multi_coo_matrix& other_matrix) :
        Parent((const Parent&)other_matrix) {}
};

template <typename IndexType, typename ValueType, class MemorySpace, typename MultMethodType,
         typename PartitionerType = cusp::graph::patoh, typename MethodType = cusp::graph::rowwise>
class multi_csr_matrix : public multi_matrix<IndexType,ValueType,MemorySpace,cusp::multi_csr_format,MultMethodType,PartitionerType,MethodType>
{
    typedef multi_matrix<IndexType,ValueType,MemorySpace,cusp::multi_csr_format,MultMethodType,PartitionerType,MethodType> Parent;

public:
    typedef MultMethodType multi_mult_method;
    typedef MethodType part_method;

    multi_csr_matrix(void) : Parent() {}

    template <typename ArrayType>
    multi_csr_matrix(size_t num_rows, size_t num_cols, size_t num_entries, const ArrayType& part)
        : Parent(num_rows, num_cols, num_entries, num_entries, part) {}

    template <typename MatrixType>
    multi_csr_matrix(const MatrixType& matrix)
        : Parent(matrix) {}

    template <typename MatrixType, typename ArrayType>
    multi_csr_matrix(const MatrixType& matrix, const ArrayType& part)
        : Parent(matrix, part) {}

    multi_csr_matrix(const multi_csr_matrix& other_matrix) :
        Parent((const Parent&)other_matrix) {}
};

template <typename IndexType, typename ValueType, class MemorySpace, typename MultMethodType,
         typename PartitionerType = cusp::graph::patoh, typename MethodType = cusp::graph::rowwise>
class multi_dia_matrix : public multi_matrix<IndexType,ValueType,MemorySpace,cusp::multi_dia_format,MultMethodType,PartitionerType,MethodType>
{
    typedef multi_matrix<IndexType,ValueType,MemorySpace,cusp::multi_dia_format,MultMethodType,PartitionerType,MethodType> Parent;

public:
    typedef MultMethodType multi_mult_method;
    typedef MethodType part_method;

    multi_dia_matrix(void) : Parent() {}

    template <typename ArrayType>
    multi_dia_matrix(size_t num_rows, size_t num_cols, size_t num_entries, const ArrayType& part)
        : Parent(num_rows, num_cols, num_entries, num_entries, part) {}

    template <typename MatrixType>
    multi_dia_matrix(const MatrixType& matrix)
        : Parent(matrix) {}

    template <typename MatrixType, typename ArrayType>
    multi_dia_matrix(const MatrixType& matrix, const ArrayType& part)
        : Parent(matrix, part) {}

    multi_dia_matrix(const multi_dia_matrix& other_matrix) :
        Parent((const Parent&)other_matrix) {}
};

template <typename IndexType, typename ValueType, class MemorySpace, typename MultMethodType,
         typename PartitionerType = cusp::graph::patoh, typename MethodType = cusp::graph::rowwise>
class multi_ell_matrix : public multi_matrix<IndexType,ValueType,MemorySpace,cusp::multi_ell_format,MultMethodType,PartitionerType,MethodType>
{
    typedef multi_matrix<IndexType,ValueType,MemorySpace,cusp::multi_ell_format,MultMethodType,PartitionerType,MethodType> Parent;

public:
    typedef MultMethodType multi_mult_method;
    typedef MethodType part_method;

    multi_ell_matrix(void) : Parent() {}

    template <typename ArrayType>
    multi_ell_matrix(size_t num_rows, size_t num_cols, size_t num_entries, const ArrayType& part)
        : Parent(num_rows, num_cols, num_entries, num_entries, part) {}

    template <typename MatrixType>
    multi_ell_matrix(const MatrixType& matrix)
        : Parent(matrix) {}

    template <typename MatrixType, typename ArrayType>
    multi_ell_matrix(const MatrixType& matrix, const ArrayType& part)
        : Parent(matrix, part) {}

    multi_ell_matrix(const multi_ell_matrix& other_matrix) :
        Parent((const Parent&)other_matrix) {}
};

template <typename IndexType, typename ValueType, class MemorySpace, typename MultMethodType,
         typename PartitionerType = cusp::graph::patoh, typename MethodType = cusp::graph::rowwise>
class multi_hyb_matrix : public multi_matrix<IndexType,ValueType,MemorySpace,cusp::multi_hyb_format,MultMethodType,PartitionerType,MethodType>
{
    typedef multi_matrix<IndexType,ValueType,MemorySpace,cusp::multi_hyb_format,MultMethodType,PartitionerType,MethodType> Parent;

public:
    typedef MultMethodType multi_mult_method;
    typedef MethodType part_method;

    multi_hyb_matrix(void) : Parent() {}

    template <typename ArrayType>
    multi_hyb_matrix(size_t num_rows, size_t num_cols, size_t num_entries, const ArrayType& part)
        : Parent(num_rows, num_cols, num_entries, num_entries, part) {}

    template <typename MatrixType>
    multi_hyb_matrix(const MatrixType& matrix)
        : Parent(matrix) {}

    template <typename MatrixType, typename ArrayType>
    multi_hyb_matrix(const MatrixType& matrix, const ArrayType * part)
        : Parent(matrix, part) {}

    template <typename MatrixType, typename ArrayType>
    multi_hyb_matrix(const MatrixType& matrix, const ArrayType& part)
        : Parent(matrix, part) {}

    multi_hyb_matrix(const multi_hyb_matrix& other_matrix) :
        Parent((const Parent&)other_matrix) {}
};

/*! \}
 */

} // end namespace cusp

